from fastapi import FastAPI, File, UploadFile, BackgroundTasks
from fastapi.responses import JSONResponse
from feature_extractor import FeatureExtractor
from faiss_db import FaissDatabase
import uuid
import os
import torch
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

app = FastAPI()
extractor = FeatureExtractor()
db = FaissDatabase()

@app.post("/index")
async def create_index(background_tasks: BackgroundTasks, paths: list[str]):
    task_id = str(uuid.uuid4())
    background_tasks.add_task(batch_process_images, paths)
    return {"task_id": task_id, "status": "processing"}

def batch_process_images(paths):
    for path in paths:
        if not os.path.exists(path):
            continue
        feature = extractor.extract(path)
        if feature is not None:
            db.add_feature(feature, path)

@app.post("/search")
async def search_image(
    file: UploadFile = File(...),
    threshold: float = 0.7,
    top_k: int = 10
):

    # 指定目录路径
    directory_path = os.path.join(os.getcwd(), 'tmp')
    temp_path=os.path.join(directory_path,file.filename)
    # 确保目录存在
    if not os.path.exists(directory_path):
        os.makedirs(directory_path)
    
    with open(temp_path, "wb") as tmp:
        tmp.write(await file.read())
    
    # 提取特征
    query_feature = extractor.extract(temp_path)
    if query_feature is None:
        return JSONResponse(
            {"error": "Failed to process image"}, 
            status_code=400
        )
    
    # 执行搜索
    results = db.search(query_feature, top_k, threshold)
    return {"results": results}

@app.get("/stats")
async def get_status():
    return {
        "index_size": db.get_index_size(),
        "gpu_available": torch.cuda.is_available(),
        "system_status": "active"
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)